from fastapi import HTTPException, status
from tortoise.exceptions import IntegrityError

from app.crud.studies import get_medical_study_by_id
from app.models.plans import StudyPlan, StudyPlanStatus
from app.schemas.plans import (
    StudyPlanCreateSchema,
    StudyPlanUpdateSchema,
)


async def create_study_plan(study_plan_schema: StudyPlanCreateSchema) -> StudyPlan:
    study_plan_schema_dict = study_plan_schema.model_dump(exclude_unset=True)
    study = await get_medical_study_by_id(study_plan_schema_dict.pop("study_id"))
    new_study_plan = StudyPlan(study=study, **study_plan_schema_dict)
    try:
        await new_study_plan.save()
    except IntegrityError:
        raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail="Study plan already exists")
    await new_study_plan.fetch_related("study", "study__modality")
    return new_study_plan


async def get_study_plan_by_id(study_plan_id: int) -> StudyPlan:
    study_plan = await StudyPlan.filter(id=study_plan_id).first().prefetch_related("study", "study__modality")
    if not study_plan:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Study plan not found")
    return study_plan


async def update_study_plan(study_plan_update_schema: StudyPlanUpdateSchema) -> StudyPlan:
    study_plan_update_dict = study_plan_update_schema.model_dump(exclude_unset=True)
    study_plan = await get_study_plan_by_id(study_plan_update_dict.pop("id"))
    await study_plan.update_from_dict(study_plan_update_dict)
    await study_plan.save()
    return study_plan


async def delete_study_plan(study_plan_id: int) -> None:
    deleted_count = await StudyPlan.filter(id=study_plan_id).delete()
    if not deleted_count:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Study plan not found")


async def get_study_plans_by_year_and_month(year: int, month: int) -> list[StudyPlan]:
    return await StudyPlan.filter(year=year, month=month).prefetch_related("study", "study__modality")


async def get_all_study_plans_by_year(year: int) -> list[StudyPlan]:
    return await StudyPlan.filter(year=year).prefetch_related("study", "study__modality")


async def get_all_study_plans_by_period(
    start_year: int, start_month: int, end_year: int, end_month: int
) -> list[StudyPlan]:
    return await StudyPlan.filter(
        year__gte=start_year, year__lte=end_year, month__gte=start_month, month__lte=end_month
    ).prefetch_related("study", "study__modality")


async def approve_study_plan(study_plan_id: int) -> StudyPlan:
    study_plan = await get_study_plan_by_id(study_plan_id)
    study_plan.status = StudyPlanStatus.APPROVED
    await study_plan.save()
    return study_plan


async def approve_study_plans_by_ids(study_plan_ids: list[int]) -> list[StudyPlan]:
    await StudyPlan.filter(id__in=study_plan_ids).update(status=StudyPlanStatus.APPROVED)
    return await StudyPlan.filter(id__in=study_plan_ids).prefetch_related("study", "study__modality")


async def complete_study_plan(study_plan_id: int) -> StudyPlan:
    study_plan = await get_study_plan_by_id(study_plan_id)
    study_plan.status = StudyPlanStatus.COMPLETED
    await study_plan.save()
    return study_plan


async def complete_study_plans_by_ids(study_plan_ids: list[int]) -> list[StudyPlan]:
    await StudyPlan.filter(id__in=study_plan_ids).update(status=StudyPlanStatus.COMPLETED)
    return await StudyPlan.filter(id__in=study_plan_ids).prefetch_related("study", "study__modality")
